{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../learner/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from abc import ABC\n",
    "from functools import partial\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd.function import Function\n",
    "\n",
    "from torch_learner import Learner, set_seed, Metrics, LrScheduler, Callback\n",
    "\n",
    "\n",
    "set_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# center loss\n",
    "class CenterLoss(nn.Module, ABC):\n",
    "    \"\"\"使用parameter封装训练center.\n",
    "    center_loss = ||feature_embedding - center_embedding||**2\n",
    "    \"\"\"\n",
    "    def __init__(self, cls_num, feature_dim):\n",
    "        super().__init__()\n",
    "        self.cls_num = cls_num\n",
    "        self.featur_num = feature_dim\n",
    "        self.center = nn.Parameter(torch.randn(cls_num, feature_dim))\n",
    "        self.is_first = True\n",
    "\n",
    "    def forward(self, xs, ys):\n",
    "        \"\"\"\n",
    "        if self.is_first:\n",
    "            self.center = self.center.to(xs.device)\n",
    "            self.is_first = False\n",
    "        \"\"\"\n",
    "        center_exp = self.center.index_select(dim=0, index=ys.long())\n",
    "        count = torch.histc(ys.float(), bins=self.cls_num, min=0, max=self.cls_num - 1)\n",
    "        count_dis = count.index_select(dim=0, index=ys.long()) + 1\n",
    "        loss = torch.sum(torch.sum((xs - center_exp) ** 2, dim=1) / 2.0 / count_dis.float())\n",
    "        return loss\n",
    "\n",
    "\n",
    "class MyLoss(nn.Module, ABC):\n",
    "    def __init__(self, center_loss_weight: float = 1.0, cls_num: int = 10, feature_dim: int = 2):\n",
    "        super(MyLoss, self).__init__()\n",
    "        self.center_loss_func = CenterLoss(cls_num, feature_dim)\n",
    "        self.nllloss_func = nn.NLLLoss()\n",
    "        self.center_loss_weight = center_loss_weight\n",
    "\n",
    "    def forward(self, outputs, targets):\n",
    "        return self.nllloss_func(outputs[1], targets) + \\\n",
    "            self.center_loss_weight * self.center_loss_func(outputs[0], targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize callback\n",
    "class TrainCallback(Callback):\n",
    "    def __init__(self):\n",
    "        self.embeddings = []    # 2维的embedding，用于可视化\n",
    "        self.labels = []\n",
    "\n",
    "    def on_epoch_begin(self, info: dict):\n",
    "        # 每次epoch开始时，清空之前的数据\n",
    "        self.embeddings.clear()\n",
    "        self.labels.clear()\n",
    "\n",
    "    def on_loss_begin(self, info: dict):\n",
    "        # 获取每个batch得到的向量\n",
    "        self.embeddings.append(info[\"outputs\"][0])\n",
    "        self.labels.append(info[\"y\"][0].int())\n",
    "\n",
    "    def on_metric_begin(self, info: dict):\n",
    "        info[\"orig_outputs\"] = info[\"outputs\"]\n",
    "        info[\"outputs\"] = info[\"outputs\"][1]\n",
    "\n",
    "    def on_epoch_end(self, info: dict):\n",
    "        # 每个epoch结束时，对向量进行可视化\n",
    "        embeddings = torch.cat(self.embeddings, dim=0)\n",
    "        labels = torch.cat(self.labels, dim=0)\n",
    "        self.visualize(embeddings.detach().cpu().numpy(), labels.detach().cpu().numpy(), info[\"epoch\"])\n",
    "\n",
    "    @staticmethod\n",
    "    def visualize(feat, labels, epoch):\n",
    "        plt.ion()\n",
    "        c = ['#ff0000', '#ffff00', '#00ff00', '#00ffff', '#0000ff',\n",
    "             '#ff00ff', '#990000', '#999900', '#009900', '#009999']\n",
    "        plt.clf()\n",
    "        for i in range(10):\n",
    "            plt.plot(feat[labels == i, 0], feat[labels == i, 1], '.', c=c[i])\n",
    "        plt.legend(['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'], loc='upper right')\n",
    "        plt.xlim(xmin=-10, xmax=10)\n",
    "        plt.ylim(ymin=-10, ymax=10)\n",
    "        plt.text(-7.8, 7.3, \"epoch=%d\" % epoch)\n",
    "        plt.savefig('/tmp/epoch=%d.jpg' % epoch)\n",
    "        plt.draw()\n",
    "        plt.pause(0.001)\n",
    "\n",
    "\n",
    "class ValidCallback(Callback):\n",
    "    def on_metric_begin(self, info: dict):\n",
    "        info[\"orig_outputs\"] = info[\"outputs\"]\n",
    "        info[\"outputs\"] = info[\"outputs\"][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model\n",
    "\n",
    "class Model(nn.Module, ABC):\n",
    "    \"\"\"LeNet with prelu.\"\"\"\n",
    "    def __init__(self, cls_num: int = 10):\n",
    "        super(Model, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 6, 5)\n",
    "        self.relu1 = nn.PReLU()\n",
    "        self.pool1 = nn.MaxPool2d(2)\n",
    "        self.conv2 = nn.Conv2d(6, 16, 5)\n",
    "        self.relu2 = nn.PReLU()\n",
    "        self.pool2 = nn.MaxPool2d(2)\n",
    "        self.fc1 = nn.Linear(256, 120)\n",
    "        self.relu3 = nn.PReLU()\n",
    "        self.fc2 = nn.Linear(120, 84)\n",
    "        self.relu4 = nn.PReLU()\n",
    "        self.fc3 = nn.Linear(84, 2)\n",
    "        self.relu5 = nn.PReLU()\n",
    "        self.fc = nn.Linear(2, cls_num)\n",
    "        self.log_softmax = nn.LogSoftmax(dim=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.conv1(x)\n",
    "        y = self.relu1(y)\n",
    "        y = self.pool1(y)\n",
    "        y = self.conv2(y)\n",
    "        y = self.relu2(y)\n",
    "        y = self.pool2(y)\n",
    "        y = y.view(y.shape[0], -1)\n",
    "        y = self.fc1(y)\n",
    "        y = self.relu3(y)\n",
    "        y = self.fc2(y)\n",
    "        y = self.relu4(y)\n",
    "        y = self.fc3(y)\n",
    "        logit = self.relu5(y)\n",
    "        y = self.fc(logit)\n",
    "        return logit, self.log_softmax(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data\n",
    "# 1. 数据集\n",
    "train_dataset = datasets.MNIST(\n",
    "    '~/.cache/datasets/mnist',\n",
    "    download=True,\n",
    "    train=True,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))]))\n",
    "valid_dataset = datasets.MNIST(\n",
    "    '~/.cache/datasets/mnist',\n",
    "    download=True,\n",
    "    train=False,\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = Model()\n",
    "learner = Learner(model=model,\n",
    "                  train_ds=train_dataset,\n",
    "                  valid_ds=valid_dataset,\n",
    "                  metrics={'accuracy': partial(Metrics.accuracy_score, just_score=False),\n",
    "                           'precision_macro': partial(Metrics.precision_score, average='macro', just_score=False),\n",
    "                           'recall_macro': partial(Metrics.recall_score, average='macro', just_score=False),\n",
    "                           'f1_macro': partial(Metrics.f1_score, average='macro', just_score=False)},\n",
    "                  loss_func=MyLoss(),\n",
    "                  optim_func=torch.optim.Adam,\n",
    "                  batch_size=128,\n",
    "                  lr=0.001,\n",
    "                  lr_scheduler=LrScheduler.get_step_lr(step_size=5, gamma=0.2),\n",
    "                  valid_batch=-1,\n",
    "                  device=torch.device('cuda'),\n",
    "                  train_callbacks=[TrainCallback()],\n",
    "                  valid_callbacks=[ValidCallback()])\n",
    "learner.train(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
